from utils.preprocess import GlassFucker
import os

from utils.plotting import plot_img
from models.ssd.ssd import build_ssd
import torch
import torch.backends.cudnn as cudnn
import cv2
from matplotlib import pyplot as plt

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


resume = './saved_weights/test1/bestweights-[0.6569].pth'

img_path = '/media/hszc/data1/glass_data/img_no_crayon/20180612144201_1.bmp'
mask_path = '/media/hszc/data1/glass_data/mask_no_crayon/20180612144201_1.jpg'

GF = GlassFucker(resize=None)


print img_path

GF.imread(img_path, mask_path)
GF.find_bound()

pred_img = GF.img.copy()
gt_img = GF.img.copy()

print GF.img.shape
print GF.mask.shape
print GF.xmin, GF.xmax, GF.ymin, GF.ymax


patches, mask_patches, coords, bboxs = GF.grid_patch(patch_size=500, grid_num=(10, 10), check=True)


num_classes = 2
ssd_net = build_ssd('test', 300, num_classes)  # return 'nn.Module', forward return [loc pred, conf pred ,prior]

net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True  # set to False when inputsize changes

if resume:
    print('Resuming training, loading {}...'.format(resume))
    net.load_state_dict(torch.load(resume, map_location=lambda storage, loc: storage))
    print('finish loading')

net = net.cuda()

for pidx in range(patches.shape[0]):

    input_img = cv2.resize(patches[pidx], (300, 300))
    input_img, img_true, img_pred, preds, true_box = plot_img(net, input_img, annos=bboxs[pidx], transforms=None, scale =500./300)

    for i in range(len(preds)):
        print 'preds', preds[i]
        preds[i][0] += coords[i][0]
        preds[i][1] += coords[i][1]
        preds[i][2] += coords[i][0]
        preds[i][3] += coords[i][1]

        cv2.rectangle(pred_img, (preds[i][0], preds[i][1]), (preds[i][2], preds[i][3]), (230, 25, 75), thickness=20)

    for i in range(len(true_box)):
        true_box[i][0] += coords[i][0]
        true_box[i][1] += coords[i][1]
        true_box[i][2] += coords[i][0]
        true_box[i][3] += coords[i][1]
        cv2.rectangle(gt_img, (true_box[i][0], true_box[i][1]), (true_box[i][2], true_box[i][3]), (230, 25, 75), thickness=20)

    # break


plt.figure()
plt.subplot(131)
plt.imshow(GF.img)

plt.subplot(132)
plt.imshow(gt_img)

plt.subplot(133)
plt.imshow(pred_img)

plt.show()
